"""
中心词 + 上下文
"""
import re

import torch

from dataset.vocab import Vocab
from random import randint

sentence_str = "Where is the captain of china?the captain of china is beijing."
tokens = [re.sub(r"[^a-zA-Z]+", " ", sentence_str).lower().strip().split()]

vocab = Vocab(tokens, 0)
token_ids = [vocab.to_idx(line) for line in tokens]
idx_list = list(vocab.token_to_idx.values())


def _generate_center_and_context(token_ids, window_size=2):
    data = []
    # 获取每一个词
    for line in token_ids:
        for i, idx in enumerate(line):
            start = max(0, i - window_size)
            end = min(len(line), i + window_size + 1)
            context = line[start:i] + line[i + 1:end]
            for con in context:
                data.append([idx, con, 1])  # 1证明数据为正样本
            # 负样本采集
            neg_list = []
            for i in range(randint(2, 20)):
                # TODO： 实际中需要根据词频来优先选取词频较大的负样本
                j = randint(0, len(vocab) - 1)
                if idx_list[j] not in context and idx_list[j] not in neg_list:
                    neg_list.append(idx_list[j])
                    data.append([idx, idx_list[j], 0])  # 0证明数据为负样本

    return torch.tensor(data, dtype=torch.long)


data = _generate_center_and_context(token_ids)
print(data)
